# 用一个类定义一个节点
from typing import Any

from langchain_openai import ChatOpenAI
from langgraph.graph import Graph

from rich.console import Console
console = Console()

from rich.markdown import Markdown

# 初始化模型
def qw_model():
    return ChatOpenAI(
        model='qwen-max-0919',
        base_url='https://dashscope.aliyuncs.com/compatible-mode/v1',
        api_key='sk-debee146f82244268914dd2e3d98761b',
        temperature=0.7,
        top_p=0.9
    )

# 1.先定义节点的功能类(逻辑)
class AskNode:
    def __init__(self) -> None:
        self.llm = qw_model()

    def __call__(self, state) -> Any:
        resp = self.llm.invoke(state['user_input'])
        return {
            "user_input": state['user_input'],
            "ai_response": resp.content
        }

class AnswerNode:
    def __call__(self, state) -> Any:
        print(state)
        return state


# 2.构建图
workflow = Graph()
# 添加节点
workflow.add_node("node_1", AskNode())
workflow.add_node("node_2", AnswerNode())
# 添加边
workflow.add_edge("node_1", "node_2")

# 定义入口节点和出口节点
workflow.set_entry_point("node_1")
workflow.set_finish_point("node_2")

# 3.对图进行编译
app = workflow.compile()
# 调用方式
for output in app.stream({"user_input": "介绍李白"}):
    for k, v in output.items():
        # print(f"{k=},{v=}")

        # 美化输出格式
        console.print(f"node:{k}", style="white on blue")
        console.print("\n")
        console.print(Markdown(v['ai_response']))
        console.print("--"*30)
